# 导入日志模块
import logging

# 从http.server模块导入HTTPServer和BaseHTTPRequestHandler类
from http.server import HTTPServer, BaseHTTPRequestHandler

# 导入json模块
import json

# 导入secrets模块用于生成安全的随机数
import secrets
import base64
import hashlib

# 导入datetime相关类
from datetime import datetime, timedelta, timezone
from urllib.parse import parse_qs, urlparse

# 配置日志记录级别为INFO
logging.basicConfig(level=logging.INFO)
# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


# 定义OAuthServer类，作为OAuth 2.1授权服务器
class OAuthServer:
    """OAuth 2.1授权服务器"""

    # 初始化方法
    def __init__(self):
        self.clients = {}
        self.authorization_code = {}
        self.server_secret = "server_secret_2025"
        self.access_tokens = {}
        # 服务器元数据
        self.metadata = {
            # 发行者
            "issuer": "http://localhost:8001",
            # 授权端点
            "authorization_endpoint": "http://localhost:8001/oauth/authorize",
            # 令牌端点
            "token_endpoint": "http://localhost:8001/oauth/token",
            # 客户端注册端点
            "registration_endpoint": "http://localhost:8001/oauth/register",
            # 支持的作用域
            "scopes_supported": ["mcp:tools", "mcp:resources", "mcp:prompts"],
            # 支持的响应类型
            "response_types_supported": ["code"],
            # 支持的授权类型
            "grant_types_supported": ["authorization_code", "refresh_token"],
            # 令牌端点支持的认证方法
            "token_endpoint_auth_methods_supported": [
                "client_secret_basic",
                "client_secret_post",
            ],
            # 支持的PKCE方法
            "code_challenge_methods_supported": ["S256"],
            # 服务文档地址
            "service_documentation": "https://modelcontextprotocol.io/specification/2025-06-18/basic/authorization",
        }

    def create_authorization_code(
        self, client_id, redirect_uri, scope, state, code_challenge
    ):
        if client_id not in self.clients:
            return None
        code = str(secrets.token_urlsafe(32))
        self.authorization_code[code] = {
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "scope": scope,
            "state": state,
            "code_challenge": code_challenge,
            "expires_at": datetime.now(timezone.utc) + timedelta(minutes=10),
            "created_at": datetime.now(timezone.utc),
        }
        return code

    def validate_authorization_code(self, code, client_id, code_verifier):
        if code not in self.authorization_code:
            return None
        auth_code = self.authorization_code[code]
        if datetime.now(timezone.utc) > auth_code["expires_at"]:
            del self.authorization_code[code]
            return None
        if auth_code["client_id"] != client_id:
            return None
        if auth_code["code_challenge"] and code_verifier:
            code_verifier_b64 = (
                base64.urlsafe_b64encode(
                    hashlib.sha256(code_verifier.encode()).digest()
                )
                .decode()
                .rstrip("=")
            )
            if code_verifier_b64 != auth_code["code_challenge"]:
                return None
        result = auth_code.copy()
        # 授权码是一次性的，用完就销毁了
        del self.authorization_code[code]
        return result

    def create_access_token(self, client_id, scope, user_id="anonymous"):
        import jwt
        from datetime import datetime, timedelta

        # 构造JWT载荷
        payload = {
            "iss": self.metadata["issuer"],  # 发行者
            "aud": "http://localhost:8000",  # 目标资源（MCP服务器）
            "sub": user_id,
            "client_id": client_id,
            "scope": scope,
            "iat": datetime.now(timezone.utc),
            "exp": datetime.now(timezone.utc) + timedelta(hours=1),
        }
        access_token = jwt.encode(payload, self.server_secret, algorithm="HS256")
        self.access_tokens[access_token] = {
            "client_id": client_id,
            "scope": scope,
            "user_id": user_id,
            "created_at": datetime.now(timezone.utc),
            "expires_at": datetime.now(timezone.utc) + timedelta(hours=1),
        }
        return access_token

    def register_client(self, client_metadata):
        client_id = str(secrets.token_urlsafe(16))
        client_secret = str(secrets.token_urlsafe(32))
        client_info = {
            "client_id": client_id,
            "client_secret": client_secret,
            "client_secret_expires_at": 0,
            "client_id_issued_at": int(datetime.now(timezone.utc).timestamp()),
            **client_metadata,
        }
        self.clients[client_id] = client_info
        logger.info(f"新的客户端注册成功:{client_id}")
        return client_info


# 定义OAuthHandler类，继承自BaseHTTPRequestHandler
class OAuthHandler(BaseHTTPRequestHandler):
    """OAuth HTTP处理器"""

    # 初始化方法，接收oauth_server参数
    def __init__(self, *args, oauth_server=None, **kwargs):
        # 保存OAuthServer实例
        self.oauth_server = oauth_server
        # 调用父类初始化方法
        super().__init__(*args, **kwargs)

    def _handle_client_registration(self):
        try:
            # 获取请求体长度
            content_length = int(self.headers.get("Content-Length", 0))
            # 读取请求体内容
            body = self.rfile.read(content_length)
            # 解析JSON格式的客户端元数据
            client_metadata = json.loads(body.decode("utf-8"))

            # 调用OAuthServer的register_client方法注册客户端
            client_info = self.oauth_server.register_client(client_metadata)

            # 发送201响应，表示创建成功
            self.send_response(201)
            # 设置响应头为application/json
            self.send_header("Content-Type", "application/json")
            # 结束响应头
            self.end_headers()
            # 返回注册成功的客户端信息（JSON格式）
            self.wfile.write(json.dumps(client_info).encode("utf-8"))
        except Exception as e:
            logger.error(f"客户端注册错误:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def _build_query_string(self, params):
        return "&".join([f"{k}={v}" for k, v in params.items()])

    def _handle_authorization(self):
        try:
            query = urlparse(self.path).query
            params = parse_qs(query)
            client_id = params.get("client_id", [""])[0]
            redirect_uri = params.get("redirect_uri", [""])[0]
            scope = params.get("scope", [""])[0]
            state = params.get("state", [""])[0]
            code_challenge = params.get("code_challenge", [""])[0]
            resource = params.get("resource", [""])[0]
            if client_id not in self.oauth_server.clients:
                self.send_error(400, "Invalid client_id")
                return
            code = self.oauth_server.create_authorization_code(
                client_id, redirect_uri, scope, state, code_challenge
            )
            if not code:
                self.send_error(400, "Failed to create authorization code")
                return
            redirect_params = {"code": code, "state": state}
            redirect_url = f"{redirect_uri}?{self._build_query_string(redirect_params)}"
            self.send_response(302)
            self.send_header("Location", redirect_url)
            self.end_headers()
        except Exception as e:
            logger.error(f"生成授权码失败:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def _handle_authorization_code_grant(self, token_request):
        code = token_request.get("code")
        client_id = token_request.get("client_id")
        code_verifier = token_request.get("code_verifier")
        auth_code = self.oauth_server.validate_authorization_code(
            code, client_id, code_verifier
        )
        print("auth_code", auth_code)
        if not auth_code:
            logger.error(f"授权码验证失败")
            self.send_error(400, "Invalid authorization code")
            return
        logger.info(f"授权码验证成功: {auth_code}")
        access_token = self.oauth_server.create_access_token(
            client_id, auth_code["scope"]
        )
        return {
            "access_token": access_token,
            "token_type": "Bearer",
            "expires_in": 3600,
            "scope": auth_code["scope"],
        }

    def _handle_token(self):
        try:
            # 获取请求体长度
            content_length = int(self.headers.get("Content-Length", 0))
            # 读取请求体内容
            body = self.rfile.read(content_length)
            token_request = json.loads(body.decode("utf-8"))
            print("token_request", token_request)
            grant_type = token_request.get("grant_type")
            if grant_type == "authorization_code":
                response = self._handle_authorization_code_grant(token_request)
            else:
                self.send_error(400, "Unsupported grant type")
                return
            if response:
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps(response).encode("utf-8"))
        except Exception as e:
            logger.error(f"令牌处理错误:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def do_GET(self):
        try:
            if self.path.startswith("/oauth/authorize"):
                self._handle_authorization()
            else:
                self.send_error(404, "Not Found")
        except Exception as e:
            logger.error(f"GET处理错误:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def do_POST(self):
        try:
            if self.path == "/oauth/register":
                self._handle_client_registration()
            elif self.path.startswith("/oauth/token"):
                self._handle_token()
            else:
                self.send_error(404, "Not Found")
        except Exception as e:
            logger.error(f"POST处理错误:{str(e)}")
            self.send_error(500, "Internal Server Error")


# 运行OAuth授权服务器的函数
def run_oauth_server(host="localhost", port=8001):
    """运行OAuth授权服务器"""
    # 创建OAuthServer实例
    oauth_server = OAuthServer()

    # 定义处理器工厂函数，用于传递oauth_server实例
    def handler_factory(*args, **kwargs):
        return OAuthHandler(*args, oauth_server=oauth_server, **kwargs)

    # 创建HTTPServer实例，绑定主机和端口
    server = HTTPServer((host, port), handler_factory)
    # 记录服务器启动信息
    logger.info(f"OAuth授权服务器运行在 http://{host}:{port}")
    # 记录支持的端点信息
    logger.info("支持的端点:")
    logger.info(
        f"  - 服务器元数据: http://{host}:{port}/.well-known/oauth-authorization-server"
    )
    logger.info(f"  - 授权端点: http://{host}:{port}/oauth/authorize")
    logger.info(f"  - 令牌端点: http://{host}:{port}/oauth/token")
    logger.info(f"  - 客户端注册: http://{host}:{port}/oauth/register")

    try:
        # 启动服务器，进入循环监听请求
        server.serve_forever()
    except KeyboardInterrupt:
        # 捕获Ctrl+C，记录服务器停止
        logger.info("OAuth授权服务器已停止")
    finally:
        # 关闭服务器
        server.server_close()


# 判断是否为主程序入口
if __name__ == "__main__":
    # 导入argparse模块，用于解析命令行参数
    import argparse

    # 创建ArgumentParser对象，设置描述信息
    parser = argparse.ArgumentParser(description="MCP OAuth授权服务器")
    # 添加--host参数，指定服务器主机，默认localhost
    parser.add_argument(
        "--host", default="localhost", help="服务器主机 (默认: localhost)"
    )
    # 添加--port参数，指定服务器端口，默认8001
    parser.add_argument(
        "--port", type=int, default=8001, help="服务器端口 (默认: 8001)"
    )

    # 解析命令行参数
    args = parser.parse_args()
    # 启动OAuth授权服务器
    run_oauth_server(args.host, args.port)
